from typing import Callable, Generic, Iterable, Iterator, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume, skip
from zkl_aiutils_datasets.processing import DelegatingIterator
from zkl_aiutils_datasets.wrapping import wrap_dataset, wrap_dataset_iterator

AnySample = TypeVar('AnySample')
AnyState = TypeVar('AnyState')
AnySrcSample = TypeVar('AnySrcSample')
AnyDstSample = TypeVar('AnyDstSample')


class StateAttachingIterator(DatasetIterator[tuple[AnySample, AnyState]], Generic[AnySample, AnyState]):
    def __init__(self, iterator: Iterator[AnySample] | Iterable[AnySample]):
        self._iterator = wrap_dataset_iterator(iterator)

    def __next__(self) -> tuple[AnySample, AnyState]:
        sample = next(self._iterator)
        state = pause(self._iterator)
        return sample, state

    def __skip__(self, samples_n: int):
        skip(self._iterator, samples_n)

    def __pause__(self) -> AnyState:
        return pause(self._iterator)


class StateDetachingIterator(DatasetIterator[AnySample], Generic[AnySample, AnyState]):
    def __init__(self, iterator: Iterator[tuple[AnySample, AnyState]] | Iterable[tuple[AnySample, AnyState]]):
        self._iterator = wrap_dataset_iterator(iterator)
        self._state = None

    def __next__(self) -> AnySample:
        sample, self._state = next(self._iterator)
        return sample

    def __skip__(self, samples_n: int):
        if samples_n == 0:
            return
        if samples_n == 1:
            next(self._iterator)
            return
        skip(self._iterator, samples_n - 1)
        skip(self._iterator, 1)

    def __pause__(self) -> AnyState | None:
        return self._state


class StateSafeTransformedIterator(DelegatingIterator[AnyDstSample], Generic[AnyDstSample, AnyState]):
    def __init__(self,
        iterator: Iterator[tuple[AnySrcSample, AnyState]] | Iterable[tuple[AnySrcSample, AnyState]],
        transform_func: Callable[[Iterator[tuple[AnySrcSample, AnyState]]], Iterator[tuple[AnyDstSample, AnyState]]],
    ):
        iterator = StateAttachingIterator(iterator)
        iterator = transform_func(iterator)
        iterator = StateDetachingIterator(iterator)
        super().__init__(iterator)


class StateSafeTransformedDataset(Dataset[AnyDstSample]):
    def __init__(self,
        dataset: Dataset[AnySrcSample],
        transform_func: Callable[[Iterator[tuple[AnySrcSample, AnyState]]], Iterator[tuple[AnyDstSample, AnyState]]]
    ):
        self._dataset = wrap_dataset(dataset)
        self._transform_func = transform_func

    def __iter__(self) -> StateSafeTransformedIterator[AnyDstSample, AnyState]:
        return StateSafeTransformedIterator(iter(self._dataset), self._transform_func)

    def __resume__(self, state: AnyState | None = None, /) -> StateSafeTransformedIterator[AnyDstSample, AnyState]:
        return StateSafeTransformedIterator(resume(self._dataset, state), self._transform_func)
